/*
 * This file is part of the openHiTLS project.
 *
 * openHiTLS is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *
 *     http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#include "app_asn1parse.h"

static BSL_ASN1_TemplateItem g_inner_items[] = {
    { .tag = (BSL_ASN1_TAG_SEQUENCE | BSL_ASN1_TAG_CONSTRUCTED), .flags = 0, .depth = 0 }, // 容器（不 HEADERONLY，进入子项）
    { .tag = BSL_ASN1_TAG_BOOLEAN,   .flags = 0, .depth = 1 }, // flag
    { .tag = BSL_ASN1_TAG_INTEGER,   .flags = 0, .depth = 1 }, // count
};
static BSL_ASN1_Template g_inner_templ = {
    .templItems = g_inner_items,
    .templNum   = sizeof(g_inner_items)/sizeof(g_inner_items[0]),
};
static const char *g_inner_names[] = { NULL, "flag", "count" };
static HITLS_Asn1ItemInfo* g_inner_childInfos[] = { NULL, NULL, NULL };
static HITLS_Asn1ItemInfo g_inner_info = {
    .name       = "InnerSeq",
    .templ      = &g_inner_templ,
    .fieldNames = g_inner_names,
    .childInfos = g_inner_childInfos,
};

static BSL_ASN1_TemplateItem g_simple_items[] = {
    { .tag = (BSL_ASN1_TAG_SEQUENCE | BSL_ASN1_TAG_CONSTRUCTED), .flags = 0, .depth = 0 }, // 容器
    { .tag = BSL_ASN1_TAG_INTEGER,   .flags = 0, .depth = 1 }, // num
    { .tag = BSL_ASN1_TAG_IA5STRING, .flags = 0, .depth = 1 }, // name
};
static BSL_ASN1_Template g_simple_templ = {
    .templItems = g_simple_items,
    .templNum   = sizeof(g_simple_items)/sizeof(g_simple_items[0]),
};
static const char *g_simple_names[] = { NULL, "num", "name" };
static HITLS_Asn1ItemInfo* g_simple_childInfos[] = { NULL, NULL, NULL };
static HITLS_Asn1ItemInfo g_itemSimplePair = {
    .name       = "SimplePair",
    .templ      = &g_simple_templ,
    .fieldNames = g_simple_names,
    .childInfos = g_simple_childInfos,
};

static BSL_ASN1_TemplateItem g_nested_items[] = {
    { .tag = (BSL_ASN1_TAG_SEQUENCE | BSL_ASN1_TAG_CONSTRUCTED), .flags = 0,                        .depth = 0 }, // 容器
    { .tag = (BSL_ASN1_TAG_SEQUENCE | BSL_ASN1_TAG_CONSTRUCTED), .flags = BSL_ASN1_FLAG_HEADERONLY, .depth = 1 }, // inner
    { .tag = BSL_ASN1_TAG_OCTETSTRING,                            .flags = 0,                        .depth = 1 }, // payload
};
static BSL_ASN1_Template g_nested_templ = {
    .templItems = g_nested_items,
    .templNum   = sizeof(g_nested_items)/sizeof(g_nested_items[0]),
};
static const char *g_nested_names[] = { NULL, "inner", "payload" };
static HITLS_Asn1ItemInfo* g_nested_childInfos[] = { NULL, &g_inner_info, NULL };
static HITLS_Asn1ItemInfo g_itemNested = {
    .name       = "Nested",
    .templ      = &g_nested_templ,
    .fieldNames = g_nested_names,
    .childInfos = g_nested_childInfos,
};

static BSL_ASN1_TemplateItem g_int_items[] = {
    { .tag = BSL_ASN1_TAG_INTEGER, .flags = 0, .depth = 0 },
};
static BSL_ASN1_Template g_int_templ = {
    .templItems = g_int_items,
    .templNum   = sizeof(g_int_items)/sizeof(g_int_items[0]),
};
static const char *g_int_names[] = { "value" };
static HITLS_Asn1ItemInfo* g_int_childInfos[] = { NULL };
static HITLS_Asn1ItemInfo g_itemJustInteger = {
    .name       = "JustInteger",
    .templ      = &g_int_templ,
    .fieldNames = g_int_names,
    .childInfos = g_int_childInfos,
};

static HITLS_Asn1ItemInfo* g_itemInfos[] = {
    &g_itemSimplePair,
    &g_itemNested,
    &g_itemJustInteger,
    NULL
};

static int IsSeqOrSetOf(const BSL_ASN1_TemplateItem *it) {
    uint8_t pt = (uint8_t)((it->tag) & 0x1F);
    return ((it->flags & BSL_ASN1_FLAG_SAME) != 0) &&
           (pt == BSL_ASN1_TAG_SEQUENCE || pt == BSL_ASN1_TAG_SET);
}

static int NeedFallbackField(const BSL_ASN1_TemplateItem *it) {
    if (it->tag == BSL_ASN1_TAG_CHOICE || it->tag == BSL_ASN1_TAG_ANY) {
        return 1;
    }
    if (IsSeqOrSetOf(it)) {
        return 1;
    }
    return 0;
}

static int TemplateHasUnsupported(const BSL_ASN1_Template *templ) {
    if (templ == NULL || templ->templItems == NULL) {
        return 0;
    }
    for (uint32_t i = 0; i < templ->templNum; i++) {
        if (NeedFallbackField(&templ->templItems[i])) {
            return 1;
        }
    }
    return 0;
}

// encode DER length into out[], return number of bytes written (1..5)
static uint32_t EncodeDerLen(uint8_t out[5], uint32_t len) {
    if (len < 0x80) {
        out[0] = (uint8_t)len;
        return 1;
    }
    uint8_t tmp[4] = {0};
    uint32_t n = 0;
    uint32_t v = len;
    while (v > 0) {
        tmp[3 - n] = (uint8_t)(v & 0xFF);
        v >>= 8;
        n++;
    }
    out[0] = 0x80u | (uint8_t)n;
    for (uint32_t i = 0; i < n; i++) {
        out[1 + i] = tmp[4 - n + i];
    }
    return 1 + n;
}

// build a full TLV buffer from a value-only BSL_ASN1_Buffer. caller must free
static int32_t BuildTlvFromValue(const BSL_ASN1_Buffer *val, uint8_t **out, uint32_t *outLen) {
    if (val == NULL || out == NULL || outLen == NULL) {
        return HITLS_APP_INTERNAL_EXCEPTION;
    }
    if (val->buff == NULL && val->len != 0) {
        return HITLS_APP_INVALID_ARG;
    }

    uint8_t lenEnc[5];
    uint32_t lLen = EncodeDerLen(lenEnc, val->len);
    uint32_t tot = 1 + lLen + val->len;

    uint8_t *tmp = (uint8_t *)BSL_SAL_Malloc(tot);
    if (tmp == NULL) {
        return HITLS_APP_MEM_ALLOC_FAIL;
    }

    uint32_t pos = 0;
    tmp[pos++] = val->tag;
    for (uint32_t i = 0; i < lLen; i++) {
        tmp[pos++] = lenEnc[i];
    }
    if (val->len > 0) {
        errno_t err = memcpy_s(tmp + pos, tot - pos, val->buff, val->len);
        if (err != EOK) {
            BSL_SAL_Free(tmp);
            return HITLS_APP_COPY_ARGS_FAILED;
        }
        pos += val->len;
    }

    *out = tmp;
    *outLen = pos;
    return HITLS_APP_SUCCESS;
}

static inline int ItemProducesOutput(const BSL_ASN1_TemplateItem *it) {
    return ((it->tag & BSL_ASN1_TAG_CONSTRUCTED) == 0) || ((it->flags & BSL_ASN1_FLAG_HEADERONLY) != 0);
}

int32_t Asn1ParseItem(BSL_UIO *outUio, HITLS_Asn1ItemInfo *tplt, uint8_t *encode, uint32_t encLen,
                      uint32_t dlimit, uint32_t indent, uint32_t dump) {
    if (outUio == NULL || encode == NULL) {
        return BSL_NULL_INPUT;
    }

    // if no template: print a top title if provided, then fallback to dump with indent+2
    if (tplt == NULL ||
        tplt->templ == NULL ||
        tplt->templ->templItems == NULL ||
        tplt->templ->templNum == 0) {
        AppPrint(outUio, "%*s%s:\n", indent, "", tplt ? tplt->name : NULL);
        return Asn1ParseDump(outUio, encode, encLen, dlimit, indent + 2, dump);
    }

    BSL_ASN1_Template *templ = tplt->templ;

    // if template has unsupported types (CHOICE/ANY/OF), fallback to full dump.
    if (TemplateHasUnsupported(templ)) {
        AppPrint(outUio, "%*s%s:\n", indent, "", tplt->name);
        return Asn1ParseDump(outUio, encode, encLen, dlimit, indent + 2, dump);
    }

    // try to print the top title, best effort
    AppPrint(outUio, "%*s%s:\n", indent, "", tplt->name);

    uint32_t num = templ->templNum;

    BSL_ASN1_Buffer *arr = (BSL_ASN1_Buffer *)BSL_SAL_Calloc(num, sizeof(BSL_ASN1_Buffer));
    if (arr == NULL) {
        return BSL_MALLOC_FAIL;
    }

    uint8_t *p = encode;
    uint32_t left = encLen;

    // no uncertain-types handling: cb = NULL (safe because we've pre-checked unsupported)
    int32_t ret = BSL_ASN1_DecodeTemplate(templ, NULL, &p, &left, arr, num);
    if (ret != BSL_SUCCESS) {
        BSL_SAL_Free(arr);
        return ret;
    }

    // unparsed data are treated as errors
    if (left != 0) {
        BSL_SAL_Free(arr);
        return BSL_INVALID_ARG;
    }

    uint32_t outIdx = 0;
    for (uint32_t i = 0; i < num; i++) {
        const BSL_ASN1_TemplateItem *it = &templ->templItems[i];
        if (!ItemProducesOutput(it)) {
            continue;
        }

        BSL_ASN1_Buffer *val = &arr[outIdx++];
        if (val->tag == BSL_ASN1_TAG_EMPTY) {
            continue;
        }

        const char *fname = (tplt->fieldNames != NULL) ? tplt->fieldNames[i] : NULL;
        uint32_t fieldIndent = indent + 2;

        int isConstructed = ((val->tag & BSL_ASN1_TAG_CONSTRUCTED) != 0);
        HITLS_Asn1ItemInfo *child = (tplt->childInfos != NULL) ? tplt->childInfos[i] : NULL;
        int hasChildTempl = (child && child->templ && child->templ->templItems && child->templ->templNum > 0);

        if (isConstructed && hasChildTempl) {
            uint8_t *childTlv = NULL;
            uint32_t childTlvLen = 0;
            ret = Asn1ParseFormat(outUio, val, fname, dlimit, fieldIndent, dump);
            if (ret != HITLS_APP_SUCCESS) { BSL_SAL_Free(arr); return ret; }
            AppPrint(outUio, "\n", 2);
            ret = BuildTlvFromValue(val, &childTlv, &childTlvLen);
            if (ret != HITLS_APP_SUCCESS) { BSL_SAL_Free(arr); return ret; }
            HITLS_Asn1ItemInfo childNoTitle = *child; childNoTitle.name = NULL;
            ret = Asn1ParseItem(outUio, &childNoTitle, childTlv, childTlvLen, dlimit, fieldIndent + 2, dump);
            BSL_SAL_Free(childTlv);
            if (ret != BSL_SUCCESS) { BSL_SAL_Free(arr); return ret; }
        } else {
            ret = Asn1ParseFormat(outUio, val, fname, dlimit, fieldIndent + 2, dump);
            AppPrint(outUio, "\n");
            if (ret != HITLS_APP_SUCCESS) {
                BSL_SAL_Free(arr);
                return ret;
            }
        }
    }
    
    BSL_SAL_Free(arr);
    return BSL_SUCCESS;
}

HITLS_Asn1ItemInfo* LookupItemInfo(const char *name) {
    if (name == NULL || name[0] == '\0') {
        return NULL;
    }
    for (size_t i = 0; g_itemInfos[i] != NULL; i++) {
        const char *iname = g_itemInfos[i]->name;
        if (iname != NULL && strcmp(iname, name) == 0) {
            return g_itemInfos[i];
        }
    }
    return NULL;
}